import os
from openai import OpenAI
import csv
import os
from dotenv import load_dotenv
from langchain_community.document_loaders import CSVLoader
from langchain.docstore.document import Document
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_community.document_compressors.dashscope_rerank import DashScopeRerank
import pandas as pd
# from langchain.text_splitter import CharacterTextSplitter
# from langchain_community.embeddings import OpenAIEmbeddings
# from langchain_community.embeddings import HuggingFaceEmbeddings
# from langchain_community.embeddings import FastEmbedEmbeddings

from langchain_iris import IRISVector

os.environ["DASHSCOPE_API_KEY"] = "sk-ec120a9dc6ef438797c430fad512ba71"


def bailian_call():
    try:
        client = OpenAI(
            # 若没有配置环境变量，请用百炼API Key将下行替换为：api_key="sk-xxx",
            api_key=os.getenv("DASHSCOPE_API_KEY"),
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
        )

        completion = client.chat.completions.create(
            model="deepseek-r1-distill-llama-8b",  # 模型列表：https://help.aliyun.com/zh/model-studio/getting-started/models
            messages=[
                {"role": "system", "content": "你将使用 ICD-10 疾病代码表来回答问题"},
                {"role": "user", "content": "乳腺炎的ICD-10 代码是什么"},
            ],
        )
        print(completion.choices[0].message.content)
    except Exception as e:
        print(f"错误信息：{e}")


def embeddings_gen():
    try:
        client = OpenAI(
            api_key=os.getenv("DASHSCOPE_API_KEY"),
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
        )
        completion = client.embeddings.create(
            model="text-embedding-v3",
            input=[
                "风急天高猿啸哀",
                "渚清沙白鸟飞回",
                "无边落木萧萧下",
                "不尽长江滚滚来",
            ],
            encoding_format="float",
        )
        print(completion.model_dump_json())
    except Exception as e:
        print(f"错误信息：{e}")


def iris_vector(target,source):
    ## read target
    loader = CSVLoader(
        file_path=target,
        csv_args={
            "delimiter": ",",
            "quotechar": '"',
        },
        content_columns=["Disease Name"]
    )
    documents = loader.load()
    # text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
    # docs = text_splitter.split_documents(documents)

    CONNECTION_STRING = "iris://superuser:fountain@localhost:1992/DEMO"

    load_dotenv(override=True)

    embeddings = DashScopeEmbeddings(
        dashscope_api_key=os.getenv("DASHSCOPE_API_KEY"), model="text-embedding-v2"
    )
    # doc_results = embeddings.aembed_documents(documents)
    COLLECTION_NAME = os.path.splitext(os.path.basename(target))[0]

    db = IRISVector.from_documents(
        embedding=embeddings,
        documents=documents,
        collection_name=COLLECTION_NAME,
        connection_string=CONNECTION_STRING,
    )
    ## read source csv with pandas
    df_source = pd.read_csv(source)
    
    df_target = pd.read_csv(target)
    
    
    
    ## iterate each row
    for index, row in df_source.iterrows():
        ##print(row[0], row['Disease Name'])
        query = row['Disease Name']
        docsrt = db.similarity_search(query)
        ## print(docsrt)
        reranker = DashScopeRerank(model="gte-rerank")
        rerankres =reranker.rerank(documents=docsrt, query=query, top_n=1)[0]
        print(docsrt[rerankres["index"]].page_content)

def query_only(query):
    embeddings = DashScopeEmbeddings(
        dashscope_api_key=os.getenv("DASHSCOPE_API_KEY"), model="text-embedding-v2"
    )
    CONNECTION_STRING = "iris://superuser:fountain@localhost:1992/DEMO"
    COLLECTION_NAME = "disease"

    db = IRISVector.from_documents(
        embedding=embeddings,
        documents="",
        collection_name=COLLECTION_NAME,
        connection_string=CONNECTION_STRING,
    )
    docs_with_score = db.similarity_search_with_score(query)
    docsrt = db.similarity_search(query)
    index = 0
    for doc, score in docs_with_score:
        print("-" * 80)
        print("Score: ", score)
        print("index:", index)
        index += 1
        print(doc.page_content)
        print("-" * 80)

    reranker = DashScopeRerank(
        model="gte-rerank",
        # other params...
    )
    print(reranker.rerank(documents=docsrt, query=query, top_n=2))


## main
if __name__ == "__main__":
    ## http_call()
    ## sdk_con()
    ## bailian_call()
    ## embeddings_gen()
    iris_vector('/Users/jiliu/Documents/gitee/dic-map/python/disease_dict.csv','/Users/jiliu/Documents/gitee/dic-map/python/disease_dict_b.csv')
    ## query_only("尿道炎")
    ## rerank()
